from mindquantum.core.circuit import Circuit
from mindquantum.simulator import Simulator
import numpy as np

def create_state_prep():
    """ Return a QuantumCircuit that prepares the binary state 1.
    """
    circ = Circuit()
    circ += Circuit().x(1)
    return circ

def check(candidate):
    candidate_circuit = candidate()
    assert isinstance(candidate_circuit, Circuit)
    n_qubits = candidate_circuit.n_qubits
    sim = Simulator('mqvector', n_qubits)
    sim.apply_circuit(candidate_circuit)
    candidate_sv = sim.get_qs()
    target_state = '1'.zfill(n_qubits)
    index = int(target_state, 2)
    reference_sv = np.zeros(2**n_qubits, dtype=complex)
    reference_sv[index] = 1.0
    assert np.allclose(candidate_sv, reference_sv, atol=1e-10)

    print("OK")

if __name__ == '__main__':
    check(create_state_prep)